package cn.lzm.registry;

import cn.hutool.core.collection.CollectionUtil;
import cn.lzm.config.RegistryConfig;
import cn.lzm.model.ServiceMetaInfo;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.cache.CuratorCache;
import org.apache.curator.framework.recipes.cache.CuratorCacheListener;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.details.JsonInstanceSerializer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collection;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ConcurrentSkipListSet;
import java.util.stream.Collectors;

/**
 * @author lzm
 * @Date 2024/11/3
 * @description 基于zookeeper的注册中心
 */
public class ZookeeperRegistry implements Registry {

    private static final Logger log = LoggerFactory.getLogger(ZookeeperRegistry.class);

    /**
     * 根节点
     */
    private static final String ZK_ROOT_PATH = "/rpc/zk";

    private CuratorFramework client;

    private ServiceDiscovery<ServiceMetaInfo> serviceDiscovery;

    /**
     * 注册中心缓存
     */
    private static final RegistryServiceCache REGISTRY_SERVICE_CACHE = new RegistryServiceCache();

    /**
     * 正在监听的key集合
     */
    private static final Set<String> WATCHING_KEYS = new ConcurrentSkipListSet<>();

    @Override
    public void init(RegistryConfig registryConfig) {
        client = CuratorFrameworkFactory.builder()
                .connectString(registryConfig.getAddress())
                .retryPolicy(new ExponentialBackoffRetry(Math.toIntExact(30), 3))
                .build();

        serviceDiscovery = ServiceDiscoveryBuilder.builder(ServiceMetaInfo.class)
                .client(client)
                .basePath(ZK_ROOT_PATH)
                .serializer(new JsonInstanceSerializer<>(ServiceMetaInfo.class))
                .build();
        try {
            client.start();
            serviceDiscovery.start();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void register(ServiceMetaInfo metaInfo) throws Exception {
        ServiceInstance<ServiceMetaInfo> serviceInstance = buildServiceInstance(metaInfo);
        serviceDiscovery.registerService(serviceInstance);
    }

    private static ServiceInstance<ServiceMetaInfo> buildServiceInstance(ServiceMetaInfo metaInfo) throws Exception {
        String address = getAddress(metaInfo);
        String name = metaInfo.getServiceKey();
        return ServiceInstance.<ServiceMetaInfo>builder().id(address)
                .name(name)
                .address(address)
                .payload(metaInfo)
                .build();
    }

    @Override
    public void unregister(ServiceMetaInfo metaInfo) {
        try {
            serviceDiscovery.unregisterService(buildServiceInstance(metaInfo));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public List<ServiceMetaInfo> discoverService(String service) {
        List<ServiceMetaInfo> cache = REGISTRY_SERVICE_CACHE.readCache(service);
        if (CollectionUtil.isNotEmpty(cache)) {
            return cache;
        }
        Collection<ServiceInstance<ServiceMetaInfo>> serviceInstances;
        try {
            serviceInstances = serviceDiscovery.queryForInstances(service);
            List<ServiceMetaInfo> metaInfos = serviceInstances.stream().map(ServiceInstance::getPayload).collect(Collectors.toList());
            if (CollectionUtil.isEmpty(metaInfos)) {
                throw new RuntimeException(String.format("未找到%s服务节点", service));
            }
            for (ServiceMetaInfo info : metaInfos) {
                watch(info.getServiceNodeKey());
            }
            REGISTRY_SERVICE_CACHE.writeCache(service, metaInfos);
            return metaInfos;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void destroy() {
        log.info("节点已下线");
        if (client != null) {
            client.close();
        }
    }

    @Override
    public void heartBeat() {
        // 不需要自动续期，只要节点没挂掉，zk就会一直有数据。
    }

    @Override
    public void watch(String serviceNodeKey) {
        if (!WATCHING_KEYS.add(serviceNodeKey)) {
            return;
        }
        CuratorCache cache = CuratorCache.build(client, ZK_ROOT_PATH + "/" + serviceNodeKey);
        cache.start();
        cache.listenable().addListener(
                CuratorCacheListener.builder()
                        .forDeletes(childData -> {
                            String path = childData.getPath();
                            log.info("被删除的节点数据为：{}", path);
                            REGISTRY_SERVICE_CACHE.clearCache(serviceNodeKey);
                        })
                        .forChanges((oldNode, newNode) -> {
                            String path = oldNode.getPath();
                            log.info("被更新的节点数据为：{}", path);
                            REGISTRY_SERVICE_CACHE.clearCache(serviceNodeKey);
                        })
                        .build()
        );
    }

    /**
     * 获取节点的地址
     *
     * @param info 生产者的注册信息
     * @return 生产者的节点地址
     */
    private static String getAddress(ServiceMetaInfo info) {
        return info.getHost() + ":" + info.getPort();
    }


}
